#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
    @version: 0.0.1
    @author: Boxi Ro
    @license: Apache Licence
    @contact: ygroup@gmail.com
    @site: https://gitee.com/boxigg/boxquant
    @software: PyCharm
    @file: tick_stra.py
    @time: 2019/11/7

    功能说明: 基于 Bars 的策略
    使用说明
    使用例子
    """

import copy
import numpy as np
import pandas as pd

from ..indicator.df_tick import *
from ..indicator.df_bar import *
from strategy.istrategy import IStrategy
from indicator.indicators import NR

class Stra(IStrategy):
    def __init__(self, parent=None, stra_ee=None, period=0, params=None, setting=None, data_ee=None):
        super(Stra, self).__init__(parent, stra_ee, period, params, setting, data_ee)

    def on_ticks(self, ticks: Ticks, **kwargs):
        """
        基于Tick的历史数据计算相关指标及策略信号
        :param ticks:
        :return:
        """
        super().on_ticks(ticks, **kwargs)

    def on_tick(self, tick, **kwargs):
        """
        根据即时数Tick据及计算出的指标计算策略信号
        :param tick:
        :param kwargs:
        :return:
        """
        super().on_tick(tick)

    def on_bar(self, bar, **kwargs):
        """
         根据即时数Bar据及计算出的指标计算策略信号
        :param bar:
        :return:
        """
        super().on_bar(bar)

    def on_bars(self, bars, **kwargs):
        """
        基于Bar的历史数据计算相关指标及策略信号
        :param bars:
        :param kwargs:
        :return:
        """
        super().on_bars(bars, **kwargs)


def zig(bars, dp=None, dr=None, price_sgn='close', n=20):
    """

    :param bars: Bars
    :param dp: 最小价格差
    :param dr: 最小波动,
    :param price_sgn: 计算价差或波动的价格, 可取值 close,open,high,low,hl,lh,
    :param n: 计算绝对离差的周期数
    :return:pd.dataframe，包含：peak,peak_is,peak_ok,peak_pos,valley,valley_is,valley_ok,valley_pos,pv
    """
    df = copy.copy(bars)
    data_len = len(df)

    p = [0] * data_len  # peak
    pp = [0] * data_len  # peak position
    pok = [0] * data_len  # 确认 peak
    pis = [0] * data_len  #
    v = [0] * data_len  # valley
    vp = [0] * data_len  # valley position
    vok = [0] * data_len  #
    vis = [0] * data_len  #
    pok_p = [0] * data_len  # 峰确认位置
    pok_v = [0] * data_len  # 峰确认价格
    vok_p = [0] * data_len  # 谷确认位置
    vok_v = [0] * data_len  # 谷确认价格

    curPV = 0  # 1=peak,-1=valley, 0=unknown

    if price_sgn.upper() == 'HL' or price_sgn.upper() == 'LH':
        h = df['high']
        l = df['low']
    else:
        h = df[price_sgn]
        l = df[price_sgn]
    if dp is not None:
        d_p = [dp] * data_len
        d_v = [dp]*data_len
    elif dr is not None:
        d_p = h * dr
        d_v = l * dr
    else:
        nr = NR(bars,n=n )
        d_p = nr
        d_v = nr
    i = 0

    p[i] = h[i]
    v[i] = l[i]
    pp[i] = i
    vp[i] = i
    while i < data_len - 1:
        p0 = p[i]
        pp0 = pp[i]
        v0 = v[i]
        vp0 = vp[i]
        i += 1
        p[i] = p[i - 1]
        pp[i] = pp[i - 1]
        v[i] = v[i - 1]
        vp[i] = vp[i - 1]
        if curPV == 0:
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                if (p[i] - v[i]) > d_v[i]:  # 谷被确认
                    vis[vp[i]] = 1
                    vok[i] = 1
                    v[i] = l[i]
                    vp[i] = i
                    curPV = -1
                    vok_p[i] = i
                    vok_v[i] = p[i]
            if v0 > l[i]:  # 仓新低
                v[i] = l[i]
                vp[i] = i
                if (p[i] - v[i]) > d_p[i]:  # 峰被确认
                    pis[pp[i]] = 1
                    pok[i] = 1
                    p[i] = h[i]
                    pp[i] = i
                    curPV = 1
                    pok_p[i] = i
                    pok_v[i] = v[i]
        elif curPV == 1:
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                if (p[i] - v[i]) > d_v[i]:  # 谷被确认
                    vis[vp[i]] = 1
                    vok[i] = 1
                    v[i] = l[i]
                    vp[i] = i
                    curPV = -1
                    vok_p[i] = i
                    vok_v[i] = p[i]
            if v0 > l[i]:  # 仓新低
                p[i] = h[i]
                pp[i] = i
                v[i] = l[i]
                vp[i] = i
        else:  # curPV==-1
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                v[i] = l[i]
                vp[i] = i
            if v0 > l[i]:  # 仓新低
                v[i] = l[i]
                vp[i] = i
                if (p[i] - v[i]) > d_p[i]:  # 峰被确认
                    pis[pp[i]] = 1
                    pok[i] = 1
                    p[i] = h[i]
                    pp[i] = i
                    curPV = 1
                    pok_p[i] = i
                    pok_v[i] = v[i]

    peak_is = pd.Series(pis, index=df.index, name='peak_is')
    peak_ok = pd.Series(pok, index=df.index, name='peak_ok')
    peak = pd.Series(p, index=df.index, name='peak')
    peak_pos = pd.Series(pp, index=df.index, name='peak_pos')

    valley_is = pd.Series(vis, index=df.index, name='valley_is')
    valley_ok = pd.Series(vok, index=df.index, name='valley_ok')
    valley = pd.Series(v, index=df.index, name='valley')
    valley_pos = pd.Series(vp, index=df.index, name='valley_pos')
    pok_v = pd.Series(pok_v, index=df.index, name='p_o_v')
    pok_p = pd.Series(pok_p, index=df.index, name='p_o_p')
    vok_v = pd.Series(vok_v, index=df.index, name='v_o_v')
    vok_p = pd.Series(vok_p, index=df.index, name='v_o_p')
    df[peak.name] = peak
    df[peak_is.name] = peak_is
    df[peak_ok.name] = peak_ok
    df[peak_pos.name] = peak_pos
    df[valley.name] = valley
    df[valley_is.name] = valley_is
    df[valley_ok.name] = valley_ok
    df[valley_pos.name] = valley_pos
    df['peak_ok_pos'] = pok_p
    df['peak_ok_value'] = pok_v
    df['valley_ok_pos'] = vok_p
    df['valley_ok_value'] = vok_v

    pv_is = peak_is - valley_is
    pv_ok = peak_ok - valley_ok
    pv = (peak + valley) / 2
    pv_ok_value = df['peak_ok_value'] + df['valley_ok_value']
    pv_ok_pos = df['peak_ok_pos']

    df['pv_ok_value'] = pv_ok_value
    df['pv_ok_pos'] = pv_ok_pos
    df['pv_ok'] = pv_ok
    df['pv'] = pd.Series(pv, index=df.index, name='pv')
    # df['pv_is']=df.query('valley_is==1')['valley_is']
    df['pv_is'] = pd.Series(pv_is, index=df.index, name='pv_is')
    df['pv_pos'] = peak_pos
    diff_peak = peak[peak_is == 1].diff(1)
    df['d_p'] = diff_peak  # 两个相邻峰的价格差
    df['d_p_p'] = peak_pos[peak_is == 1].diff(1)  # 两个相邻峰间的周期数
    df['d_p_o_v'] = pok_v[peak_ok == 1].diff(1)   # 两个相邻峰确认时的价格差
    df['d_p_o_p'] = pok_p[peak_ok == 1].diff(1)   # 两个相邻峰确认时的周期差
    diff_valley = valley[valley_is == 1].diff(1)
    df['d_v'] = diff_valley
    df['d_v_p'] = valley_pos[valley_is == 1].diff(1)
    df['d_v_o_v'] = vok_v[valley_ok == 1].diff(1)
    df['d_v_o_p'] = vok_p[valley_ok == 1].diff(1)
    diff_pv = pv[pv_is != 0].diff(1)
    df['d_pv'] = diff_pv  # 相邻峰谷间的价格差
    df['d_pv_p'] = df[pv_is != 0]['pv_pos'].diff(1)  # 相邻峰谷间的周期数
    df['d_pv_o_v'] = pv_ok_value[pv_ok != 0].diff(1)  # 相邻峰谷间确认时的价格差
    df['d_pv_o_p'] = pv_ok_pos[pv_ok != 0].diff(1)  # 相邻峰谷确认时的周期数
    # print(pis)
    # print(peak_is)
    return df


def main():
    from data.data_server import  get_ticks_df,get_bars_df
    from strategy.clc_bt_rslt import clc_long_td_rslt,clc_short_td_rslt
    symbol = 'rb2001'
    df = get_bars_df(symbol)
    bars = DFBar(data=df)
    zg = zig(bars,dr=0.01)
    print(zg.describe())


if __name__ == '__main__':
    main()
